from openpyxl import load_workbook
from openpyxl import Workbook


class RawExcel:
    """
    一个原始的Excel原型。
    """
    def __init__(self, read_path=None):
        self.lwb = load_workbook(filename=read_path)

    def get_all_sheet_name(self):
        """
        获取所有工作簿名称，工作簿的名称就是Worksheet.title变量的值。
        """
        sheet_names = []
        for name in self.lwb:
            sheet_names.append(name.title)

        return sheet_names


class SingleExcel:
    """
    读取或者写入一个excel表格。
    support：同时对多个sheet工作簿进行写入。
    """
    def __init__(self, data_type='list', read_path=None, read_column=None, set_read_sheet=None,
                 write_path=None, start=1, end=9999):
        """
        读取一个excel需要指定：read_path, read_column, set_read_sheet。
        写入一个excel需要指定：write_path, set_write_sheet。
        start：读取表格的行头。
        end：读取表格的行尾。
        """
        self.read_path = read_path
        self.write_path = write_path
        self.write_path = write_path
        self.data = None
        self._start = start
        self._end = end
        self.writer = None
        self.reader = None
        self._titles = {}   # 保存多个sheet标题对应对象。
        self.wb = Workbook()
        if self.read_path and set_read_sheet:
            self.lwb = load_workbook(filename=self.read_path)
            self.reader = self.lwb[set_read_sheet]
            self.data = self.read_all_sheet(data_type, read_column)

    def _read_by_list(self, _read_column):
        sheets = []
        for row in range(self._start, self._end):
            row_data = []
            for column in _read_column:
                row_data.append(self.reader[column + str(row)].value)
            # 一行中没有任何数据就退出。
            if not any(row_data):
                break
            sheets.append(row_data)
        return sheets

    def _read_by_dict(self, _read_column):
        sheets = {}
        for row in range(self._start, self._end):
            row_sheets = {}
            for column in _read_column:
                c_r = column + str(row)
                v = self.reader[c_r].value
                if v:
                    row_sheets.setdefault(c_r, v)
            if row_sheets:
                sheets.update(row_sheets)
            else:
                break
        return sheets

    def read_all_sheet(self, data_type, _read_column):
        """
        指定excel内容保存成list还是dict。
        """
        if data_type == 'list':
            return self._read_by_list(_read_column)
        elif data_type == 'dict':
            return self._read_by_dict(_read_column)

    def write_all_sheet(self, sheets: list, title: str):
        """
        将二维数组写入excel，需要指定sheet的工作簿名称。
        note：在写入sheet之前需要先添加sheet（调用add_sheet方法）。
        """
        if len(self._titles) == 0:
            raise IndexError("please call add_sheet add sheet.")

        self.writer = self._titles[title]
        abc = 'a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,aa,ab,ac,' \
              'ad,ae,af,ag,ah,ai,aj,ak,al,am,an,ao,ap,aq,ar,as,at,au,av,aw,ax,ay,az,'\
              'ba,bb,bc,bd,be,bf,bg,bh,bi,bj,bk,bl,bm,bn,bo,bp,bq,br,bs,bt,bu,bv,bw,bx,by,bz,'.split(',')
        row = 1

        for line in sheets:
            for n, data in enumerate(line):
                self.writer[abc[n] + str(row)] = data
            row += 1
        self.wb.save(filename=self.write_path)

    def add_sheet(self, title: str):
        """
        添加一个新的sheet表，目的是可以同时对多个sheet进行写入。
        """
        self._titles[title] = self.wb.create_sheet(title, len(self._titles))



class NotContinuousExcel(SingleExcel):
    """
    非连续行表格。
    """
    def __init__(self, _type='list', read_path=None, set_read_sheet=None, read_column=None, write_path=None, set_write_sheet=None):
        super().__init__(_type='list', read_path=read_path, set_read_sheet=set_read_sheet, read_column=read_column, write_path=write_path, set_write_sheet=set_write_sheet)

    def read_all_sheet(self, _read_column, _type):
        """
        允许有一行空格，但连续两行是空行则停止。
        :param _read_column:
        :param _type:
        :return:
        """
        sheets = []
        jump = 0
        for row in range(1, 9999):

            row_data = []
            for column in _read_column:
                row_data.append(self.reader[column + str(row)].value)
            # 连续两行没有数据就停止。
            if not any(row_data):
                if jump == 0:
                    jump = row
                if jump + 2 == row:
                    break
                continue    # 空行就跳过改行数据的保存。
            else:
                jump = 0
            sheets.append(row_data)
        return sheets
